#include "txr_algo_dlm_cls_demo.h"
#include <QFileDialog>
#include <QTextStream>
#include <QDir>
#include <QFile>
#include <sstream>
#include <iomanip> 
#include "../txr_algo_dlm_cls/txr_algo_dlm_cls.h"
#include "../tk_utility/tk_timer.hpp"
#include <QThread>
using namespace txr_algo_dlm_cls;

#include <opencv2\highgui.hpp>
#include <opencv2\imgproc.hpp>

/**
 * @brief 读取图像文件并转换为标准格式
 * @param file_path 图像文件路径
 * @param v_data 输出的图像数据
 * @param row 图像高度
 * @param col 图像宽度
 * @return 是否读取成功
 */
bool ReadImage(QString file_path, QVector<float> &v_data, int & row, int &col)
{
	using namespace cv;
	Mat src = imread(file_path.toLocal8Bit().data(), IMREAD_ANYCOLOR | IMREAD_ANYDEPTH);
	if (src.empty())
	{
		return false;
	}
	if (src.channels() != 3 && src.channels() != 1)
	{
		return false;
	}
	row = src.rows;
	col = src.cols;

	v_data.resize(row * col * 3);

	int img_size = row*col;
	if (src.channels() == 3)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3w color_value = src.ptr<Vec3w>(i)[j];
					v_data[i*col + j] = float(color_value[2]) / 65535.f;
					v_data[img_size + i*col + j] = float(color_value[1]) / 65535.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec3b color_value = src.ptr<Vec3b>(i)[j];
					v_data[i*col + j] = float(color_value[2]) / 255.f;
					v_data[img_size + i*col + j] = float(color_value[1]) / 255.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}
	if (src.channels() == 1)
	{
		if (src.depth() == CV_16U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<ushort, 1> color_value = src.ptr<Vec<ushort, 1>>(i)[j];
					v_data[i*col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size + i*col + j] = float(color_value[0]) / 65535.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 65535.f;
				}
			}
		}
		if (src.depth() == CV_8U)
		{
			for (int i = 0; i < row; i++)
			{
				for (int j = 0; j < col; j++)
				{
					Vec<uchar, 1> color_value = src.ptr<Vec<uchar, 1>>(i)[j];
					v_data[i*col + j] = float(color_value[0]) / 255.f;
					v_data[img_size + i*col + j] = float(color_value[0]) / 255.f;
					v_data[img_size * 2 + i*col + j] = float(color_value[0]) / 255.f;
				}
			}
		}
	}

	return true;
}

/**
 * @brief 在图像上标注分类结果
 * @param unit 检测单元
 * @param dest_path 目标保存路径
 * @param info 结果信息
 */
void MarkImage(st_detect_unit unit, QString dest_path, st_result_info &info)
{
	st_img_rgb img = unit.img;
	using namespace cv;
	Mat mat;
	mat.create(img.h, img.w, CV_8UC3);
	int img_size = img.h*img.w;
	for (int i = 0; i < img.h; i++)
	{
		for (int j = 0; j < img.w; j++)
		{
			mat.at<Vec3b>(i, j)[2] = img.pp_rgbf[0][i*img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[1] = img.pp_rgbf[1][i*img.w + j] * 255.f;
			mat.at<Vec3b>(i, j)[0] = img.pp_rgbf[2][i*img.w + j] * 255.f;
		}
	}
	std::stringstream stream;
	stream << std::fixed << std::setprecision(2) << info.cls << "  p:" << info.prob;
	cv::Point pt(50, 50);
	putText(mat, stream.str(), pt, 0, 0.5, cv::Scalar(0, 0, 255), 0.5);
	
	{
		QFileInfo file_info(QDir::toNativeSeparators(dest_path));
		QDir dir;
		dir.mkpath(file_info.path());
	}
	imwrite(dest_path.toLocal8Bit().data(), mat);
}

/**
 * @brief 主窗口构造函数
 */
txr_algo_dlm_cls_demo::txr_algo_dlm_cls_demo(QWidget *parent)
    : QMainWindow(parent)
{
    ui.setupUi(this);
	connect(ui.pushButtonOpenTestDir, SIGNAL(clicked()), this, SLOT(OnOpenTest()));
	connect(ui.pushButtonOpenModelDir, SIGNAL(clicked()), this, SLOT(OnOpenModel()));
	connect(ui.pushButtonTest, SIGNAL(clicked()), this, SLOT(OnTest()));
	ui.progressBar->setRange(0, 100);
	ui.progressBar->setValue(0);
}

/**
 * @brief 打开测试图像目录
 */
void txr_algo_dlm_cls_demo::OnOpenTest()
{
	QString dir = QFileDialog::getExistingDirectory(this);
	if (dir.isEmpty())
	{
		return;
	}
	ui.labelDirTest->setText(dir);
}

/**
 * @brief 打开模型目录
 */
void txr_algo_dlm_cls_demo::OnOpenModel()
{
	QString dir_path = QFileDialog::getExistingDirectory(this);
	if (dir_path.isEmpty())
	{
		return;
	}
	ui.labelDirModel->setText(dir_path);

	QDir dir(dir_path);
	QStringList filters;
	filters << "*.cls1" << "*.cls2" << "*.ocls2";  // 添加.ocls2

	QStringList files = dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
	ui.comboBoxModel->clear();
	ui.comboBoxModel->addItems(files);
}

/**
 * @brief 开始测试
 */
void txr_algo_dlm_cls_demo::OnTest()
{
	st_test_info info;

	info.test_dir = ui.labelDirTest->text();

	info.model_path = ui.labelDirModel->text() + "\\" + ui.comboBoxModel->currentText();

	info.save_img = ui.checkBoxSaveImg->isChecked();

	DlmTest(info);
}

/**
 * @brief 执行分类测试
 * @param info 测试参数
 */
void txr_algo_dlm_cls_demo::DlmTest(st_test_info &info)
{
	using namespace txr_algo_dlm_cls;
	TxrAlgoDlmCls dlm;
	
	// 加载分类模型
	{
		st_dlm_cfg cfg;

		// 解析模型文件信息
		QFileInfo model_file = QFileInfo(info.model_path);
		QString model_name = model_file.absoluteDir().path() + "\\" + model_file.baseName();
		QString model_suffix = model_file.suffix();

		// 设置模型路径
		sprintf_s(cfg.file_path_name, MAX_PATH_BYTES, "%s", model_name.toLocal8Bit().data());

		// 根据模型后缀选择对应的DLL类型
		// cls1 -> txr_algo_dlm_cls_v1.dll
		// cls2 -> txr_algo_dlm_cls_v2.dll
		// ocls2 -> txr_algo_dlm_cls_v2_o.dll
		if (model_suffix == "cls1")
		{
			cfg.type = DLM_CLS_TYPE_V1;
			// 可以在这里添加V1版本DLL的特定配置
		}
		else if (model_suffix == "cls2") 
		{
			cfg.type = DLM_CLS_TYPE_V2;
			// 可以在这里添加V2版本DLL的特定配置
		}
		else if (model_suffix == "ocls2")
		{
			cfg.type = DLM_CLS_TYPE_V2_O;  // 需要在枚举中定义这个新类型
			// 可以在这里添加V2O版本DLL的特定配置
		}
		else
		{
			// qWarning() << "Unsupported model type:" << model_suffix;
			return;
		}

		// 加载对应版本的DLL
		if (!dlm.Load(cfg))
		{
			// qWarning() << "Failed to load model:" << model_name;
			return;
		}
	}
	
	// 获取模型的batch size
	int batch_size = dlm.BatchSize();

	// 清理之前的结果目录
	{
		QDir dir(info.test_dir + "\\result");
		dir.removeRecursively();
	}
	
	// 清空结果显示区域
	ui.plainTextEditInfo->clear();

	// 获取测试图像列表
	QStringList test_files;
	{
		QDir dir(info.test_dir);
		QStringList filters;
		filters << "*.png" << "*.bmp";  // 支持png和bmp格式
		test_files =
			dir.entryList(filters, QDir::Files | QDir::Readable, QDir::Name);
	}

	// 初始化进度显示
	int file_count = 0;
	ui.progressBar->setValue(0);

	// 分配batch处理所需的缓冲区
	QVector<QString> v_img_path;        // 存储图像路径
	QVector<QVector<float>> vv_data;    // 存储图像数据
	QVector<txr_algo_dlm_cls::st_detect_unit> v_unit;  // 存储检测单元
	v_img_path.resize(batch_size);
	vv_data.resize(batch_size);
	v_unit.resize(batch_size);

	// 批量处理图像
	int batch_count = 0;
	for each (QString file in test_files)
	{
		// 填充batch数据
		if (batch_count < batch_size)
		{
			// 准备图像路径和数据缓冲区
			QString & img_path = v_img_path[batch_count];
			img_path = info.test_dir + "\\" + file;
			QVector<float> & v_data = vv_data[batch_count];
			txr_algo_dlm_cls::st_detect_unit & unit = v_unit[batch_count];
			
			// 读取图像
			int row, col;
			if (!ReadImage(img_path, v_data, row, col))
			{
				continue;
			}
			
			// 设置图像信息
			unit.img.w = col;
			unit.img.h = row;
			unit.img.c = 3;
			// 设置RGB通道数据指针
			unit.img.pp_rgbf[0] = v_data.data() + col*row * 0;  // R通道
			unit.img.pp_rgbf[1] = v_data.data() + col*row * 1;  // G通道
			unit.img.pp_rgbf[2] = v_data.data() + col*row * 2;  // B通道
			unit.param.prob_thres = 0.1;  // 设置概率阈值

			batch_count++;
		}

		// 如果batch未满且不是最后一张图，继续填充
		if (batch_count != batch_size && file != test_files.last())
		{
			continue;
		}

		// 执行推理并计时
		{
			tk::timer::universal_short_timer timer;
			dlm.Detect(v_unit.data(), batch_count);
			printf("dlm detect use time: %f ms \n", timer.elapsed_ms());
		}
		
		// 处理推理结果
		for (int i = 0; i < batch_count; ++i)
		{
			st_result_info result = v_unit[i].result;
			QFileInfo file_info = QFileInfo(v_img_path[i]);

			// 构建结果保存路径（按类别分类）
			QString mark_file = info.test_dir + QString("\\result\\%1\\").arg(result.cls) 
							  + file_info.baseName() + ".png";
			
			// 创建结果目录
			{
				QFileInfo file_info(QDir::toNativeSeparators(mark_file));
				QDir dir;
				dir.mkpath(file_info.path());
			}
			
			// 保存结果图像
			QFile::copy(v_img_path[i].toLocal8Bit(), mark_file.toLocal8Bit());
			//MarkImage(unit, mark_file, result);  // 可选的标注功能

			// 显示结果信息
			QString text = QString("no:%1 file:%2 cls:%3 prob:%4 \n")
				.arg(file_count)
				.arg(file_info.baseName())
				.arg(result.cls)
				.arg(result.prob);
			ui.plainTextEditInfo->insertPlainText(text);

			file_count++;
		}
		
		// 重置batch计数
		batch_count = 0;

		// 更新进度条
		ui.progressBar->setValue(file_count * 100 / test_files.size());
		QCoreApplication::processEvents();  // 处理UI事件，保持界面响应
	}
	
	// 完成处理
	ui.progressBar->setValue(100);
	printf("done!==========================================================\n");
}